#coding=utf8
from __future__ import division
from torch import nn
import torch
import torch.utils.data as torchdata
from torchvision import datasets,transforms
import os,time
import pandas as pd
import numpy as np
import torch.optim as optim
from torch.autograd import Variable
from torch.optim import lr_scheduler
from utils.train import train,trainlog
from sklearn.preprocessing import LabelEncoder
from torch.nn import CrossEntropyLoss
from twdata.twdataset import TWdata2
from models.xception_1d import xception
from models.Resnet_1d import resnet20_1d
from sklearn.model_selection import train_test_split
from models.LeNet import LeNet,MLPNet
import logging
from twdata.twaugment import Compose, AddNoise, RandomAmplitude, DownSample, FlowNormalize, \
                        AddAxis, CenterCrop, RandomShiftCrop


class TWAug(object):
    def __init__(self):
        self.augment = Compose([
            AddNoise(A=0.1),
            RandomAmplitude(l=0.9,h=1.1),
            RandomShiftCrop(),
            FlowNormalize(),
            AddAxis()
        ])

    def __call__(self, spct):
        return self.augment(spct)

class TWAugVal(object):
    def __init__(self):
        self.augment = Compose([
            CenterCrop(),
            FlowNormalize(),
            AddAxis()
        ])

    def __call__(self, spct):
        return self.augment(spct)




os.environ["CUDA_VISIBLE_DEVICES"] = "0"

rawdata_root = '/media/gserver/data/tianwen/rawdata'
usecuda = 1
start_epoch = 0
batch_size = 128*3
epoch_num = 130
save_inter = 5



# data prepare
train_data_root = os.path.join(rawdata_root, 'first_train_data')
train_index = pd.read_csv(os.path.join(rawdata_root, 'first_train_index_20180131.csv'))


le = LabelEncoder()
train_index['type'] = le.fit_transform(train_index['type'])
print le.classes_

index_train = pd.read_csv(os.path.join(rawdata_root,'first_train_resample.csv'))
index_val = pd.read_csv(os.path.join(rawdata_root,'first_val.csv'))
index_test = pd.read_csv(os.path.join(rawdata_root,'first_test.csv'))

class_sample_count = index_train.type.value_counts()
weights =  [int(class_sample_count.max()/class_sample_count[x]) for x in range(len(class_sample_count))]
print weights

print index_train.type.value_counts()
print index_val.type.value_counts()
print index_test.type.value_counts()




data_set = {}
data_set['train'] = TWdata2(index_pd = index_train,
                           data_root=rawdata_root,
                           classes = le.classes_,
                           transform=TWAug(),
                           )

data_set['val'] = TWdata2(index_pd = index_val,
                         data_root=rawdata_root,
                            classes = le.classes_,
                            transform=TWAugVal(),
                         )
#
data_loader = {}
data_loader['train'] = torchdata.DataLoader(data_set['train'], 32*3, num_workers=4,
                                            shuffle=True, pin_memory=True)
data_loader['val'] = torchdata.DataLoader(data_set['val'], batch_size=128*3, num_workers=4,
                                          shuffle=False, pin_memory=True)


print 'dataset: %d,%d'%(len(data_set['train']), len(data_set['val']))
#
#
#
# model prepare
resume = None
model = resnet20_1d(num_classes=4)

model = torch.nn.DataParallel(model)
if resume:
    print('resuming finetune from %s'%resume)
    model.load_state_dict(torch.load(resume))

model = model.cuda()
#

optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=1e-5)
criterion = CrossEntropyLoss()
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=100, gamma=0.5)



save_dir = '/media/gserver/models/tianwen/SMresample_resnet20_crop'
if not os.path.exists(save_dir):
    os.makedirs(save_dir)
logfile = '/media/gserver/models/tianwen/SMresample_resnet20_crop/trainlog.log'
trainlog(logfile)

best_acc,best_model_wts = train(model,
                                epoch_num,
                                batch_size,
                                start_epoch,
                                optimizer,
                                criterion,
                                exp_lr_scheduler,
                                data_set,
                                data_loader,
                                usecuda,
                                save_inter,
                                save_dir)

